//
// Created by Cybaster on 2024/3/30.
//

#include "image_stitch.h"

#include <random>
#include <opencv2/features2d.hpp>

using namespace cv;
using namespace std;


namespace stitch {

    std::vector<cv::Vec3f>
    applyH_2D(const std::vector<cv::Vec3f> &geomObjects, const cv::Matx33f &H, GeometryType type) {
        std::vector<cv::Vec3f> result;

        switch (type) {
            case GEOM_TYPE_POINT: {
                for (const auto &geomObject: geomObjects) {
                    result.push_back(H * geomObject);
                }
            }
                break;
            case GEOM_TYPE_LINE: {
                for (const auto &geomObject: geomObjects) {
                    result.push_back((H.inv()).t() * geomObject);
                }
            }
                break;
            default:
                throw std::runtime_error("Unhandled geometry type!");
        }
        return result;
    }

    cv::Matx33f getCondition2D(const std::vector<cv::Vec3f> &points2D) {
        float tx = 0.0f;
        float ty = 0.0f;
        float sx = 0.0f;
        float sy = 0.0f;
        cv::Matx33f T;

        for (const auto &i: points2D) {
            tx += i(0);
            ty += i(1);
        }
        tx /= points2D.size();
        ty /= points2D.size();

        for (const auto &i: points2D) {
            sx += abs(i(0) - tx);
            sy += abs(i(1) - ty);
        }
        sx /= points2D.size();
        sy /= points2D.size();

        T << 1 / sx, 0, -tx / sx, 0, 1 / sy, -ty / sy, 0, 0, 1;

        return T;
    }

    cv::Mat_<float> getDesignMatrix_fundamental(const std::vector<cv::Vec3f> &p1_conditioned,
                                                const std::vector<cv::Vec3f> &p2_conditioned) {
        cv::Mat_<float> A(p1_conditioned.size(), 9, CV_32F);

        for (int i = 0; i < p1_conditioned.size(); i++) {
            cv::Vec3f x1 = p1_conditioned[i];
            cv::Vec3f x2 = p2_conditioned[i];
            for (int j = 0; j < 3; j++) {
                for (int k = 0; k < 3; k++) {
                    A(i, j + 3 * k) = x1[j] * x2[k];
                }
            }
        }

        return A;
    }

    cv::Matx33f solve_dlt_fundamental(const cv::Mat_<float> &A) {
        cv::Matx33f Hhat;
        cv::Mat_<float> U, W, VT;
        cv::SVD::compute(A, W, U, VT, 4);

        for (int i = 0; i < 3; i++) {
            for (int j = 0; j < 3; j++) {
                Hhat(i, j) = VT(VT.rows - 1, i * 3 + j);
            }
        }

        return Hhat;
    }

    cv::Matx33f forceSingularity(const cv::Matx33f &F) {
        cv::Mat_<float> U, W, VT;
        cv::SVD::compute(F, W, U, VT, 4);

        cv::Mat_<float> D = cv::Mat_<float>::zeros(3, 3);
        D(0, 0) = W(0);
        D(1, 1) = W(1);

        cv::Mat_<float> F2(3, 3);
        F2 = U * D * VT;

        return F2;
    }

    cv::Matx33f decondition_fundamental(const cv::Matx33f &T1, const cv::Matx33f &T2, const cv::Matx33f &F) {
        cv::Matx33f Fu;
        Fu = T2.t() * F * T1;
        return Fu;
    }

    cv::Matx33f getFundamentalMatrix(const std::vector<cv::Vec3f> &p1, const std::vector<cv::Vec3f> &p2) {
        cv::Matx33f ConditionM_p1 = getCondition2D(p1);
        cv::Matx33f ConditionM_p2 = getCondition2D(p2);

        std::vector<cv::Vec3f> Conditioned_p1 = applyH_2D(p1, ConditionM_p1, GEOM_TYPE_POINT);
        std::vector<cv::Vec3f> Conditioned_p2 = applyH_2D(p2, ConditionM_p2, GEOM_TYPE_POINT);

        cv::Mat_<float> DesignM_Fundamental(Conditioned_p1.size(), 9, CV_32F);
        DesignM_Fundamental = getDesignMatrix_fundamental(Conditioned_p1, Conditioned_p2);

        cv::Matx33f DesignM_Estimated = solve_dlt_fundamental(DesignM_Fundamental);

        cv::Matx33f FundamentalM_Modified = forceSingularity(DesignM_Estimated);

        cv::Matx33f FundamentalM = decondition_fundamental(ConditionM_p1, ConditionM_p2, FundamentalM_Modified);

        return FundamentalM;
    }

    float getError(const cv::Vec3f &p1, const cv::Vec3f &p2, const cv::Matx33f &F) {
        Vec3f Fp1 = F * p1;
        Vec3f Fp2 = F.t() * p2;

        Vec<float, 1> errVec = ((p2.t() * F * p1) * (p2.t() * F * p1)) /
                               (pow(Fp1(0), 2) + pow(Fp1(1), 2) + pow(Fp2(0), 2) + pow(Fp2(1), 2));
        float err = errVec(0);
        return err;
    }

    float getError(const std::vector<cv::Vec3f> &p1, const std::vector<cv::Vec3f> &p2, const cv::Matx33f &F) {
        float err = 0;
        for (int i = 0; i < p1.size(); i++) {
            err += getError(p1[i], p2[i], F);
        }

        err /= p1.size();
        return err;
    }

    unsigned countInliers(const std::vector<cv::Vec3f> &p1, const std::vector<cv::Vec3f> &p2, const cv::Matx33f &F,
                          float threshold) {
        int numInliers = 0;
        for (int i = 0; i < p1.size(); i++) {

            float err = getError(p1[i], p2[i], F);
            if (err < threshold)
                numInliers++;
        }
        return numInliers;
    }

    cv::Matx33f estimateFundamentalRANSAC(const std::vector<cv::Vec3f> &p1, const std::vector<cv::Vec3f> &p2,
                                          unsigned numIterations, float threshold) {
        const unsigned subsetSize = 8;

        std::mt19937 rng;
        std::uniform_int_distribution<unsigned> uniformDist(0, p1.size() - 1);

        cv::Mat_<cv::Vec3f> my_p1(numIterations, subsetSize);
        cv::Mat_<cv::Vec3f> my_p2(numIterations, subsetSize);
        std::vector<cv::Vec3f> p1_chosen;
        std::vector<cv::Vec3f> p2_chosen;


        for (int i = 0; i < numIterations; i++) {
            for (int j = 0; j < subsetSize; j++) {
                unsigned index = uniformDist(rng);
                my_p1.at<Vec3f>(i, j) = p1[index];
                my_p2.at<Vec3f>(i, j) = p2[index];
            }
        }
        std::vector<Matx33f> my_F;
        std::vector<int> my_numInliers;
        for (int i = 0; i < numIterations; i++) {
            p1_chosen = my_p1.row(i);
            p2_chosen = my_p2.row(i);
            my_F.push_back(getFundamentalMatrix(p1_chosen, p2_chosen));
            my_numInliers.push_back(countInliers(p1, p2, my_F[i], threshold));
            p1_chosen.clear();
            p2_chosen.clear();
        }
        std::vector<int>::iterator biggest = std::max_element(my_numInliers.begin(), my_numInliers.end());
        int max_index = std::distance(my_numInliers.begin(), biggest);
        return my_F[max_index];

    }

    void visualize(const cv::Mat &img1, const cv::Mat &img2, const std::vector<cv::Vec3f> &p1,
                   const std::vector<cv::Vec3f> &p2, const cv::Matx33f &F) {
        cv::Mat img1_copy = img1.clone();
        cv::Mat img2_copy = img2.clone();

        cv::Matx33f eFundamentalMatrix = getFundamentalMatrix(p1, p2);
        std::vector<cv::Vec3f> line1s;
        std::vector<cv::Vec3f> line2s;

        cv::computeCorrespondEpilines(p1, 1, eFundamentalMatrix, line1s);
        cv::computeCorrespondEpilines(p2, 2, eFundamentalMatrix, line2s);

        for (auto p1_draw: p1) {
            cv::circle(img1_copy, cv::Point2f(p1_draw[0], p1_draw[1]), 2, cv::Scalar(0, 255, 0));
        }
        for (auto p2_draw: p2) {
            cv::circle(img2_copy, cv::Point2f(p2_draw[0], p2_draw[1]), 2, cv::Scalar(0, 255, 0));
        }

        cv::imshow("Epilines img1", img1_copy);
        cv::imshow("Epilines img2", img2_copy);

        cv::waitKey(0);
    }

    void filterMatches(const RawOrbMatches &rawOrbMatches, std::vector<cv::Vec3f> &p1, std::vector<cv::Vec3f> &p2) {

        p1.clear();
        p2.clear();

        const float ratio = 0.75f;

        for (const auto &pair: rawOrbMatches.matches_1_2) {

            if ((pair.second.closestDistance / pair.second.secondClosestDistance) > ratio) {
                continue;
            }
            auto it = rawOrbMatches.matches_2_1.find(pair.first);
            RawOrbMatches::Match value = it->second;
            if (value.closestDistance != pair.second.closestDistance) {
                continue;
            }


            p1.push_back(rawOrbMatches.keypoints1[pair.first]);
            p2.push_back(rawOrbMatches.keypoints2[pair.second.closest]);

        }
    }

    void getPointsAutomatic(const cv::Mat &img1, const cv::Mat &img2, std::vector<cv::Vec3f> &p1,
                            std::vector<cv::Vec3f> &p2) {
        RawOrbMatches result = extractRawOrbMatches(img1, img2);
        filterMatches(result, p1, p2);
    }

    void getRawPointsAutomatic(const cv::Mat &img1, const cv::Mat &img2, std::vector<cv::Vec3f> &p1,
                               std::vector<cv::Vec3f> &p2) {
        RawOrbMatches result = extractRawOrbMatches(img1, img2);
        p1 = result.keypoints1;
        p2 = result.keypoints2;
    }

    cv::Mat_<float> getDesignMatrix_homography2D(const std::vector<cv::Vec3f> &conditioned_base,
                                                 const std::vector<cv::Vec3f> &conditioned_attach) {
        cv::Mat_<float> A = cv::Mat_<float>::zeros(2 * conditioned_base.size(), 9);
        cv::Matx<float, 1, 3> A11;
        cv::Matx<float, 1, 3> A13;
        cv::Matx<float, 1, 3> A22;
        cv::Matx<float, 1, 3> A23;

        for (int i = 0; i < conditioned_base.size(); i++) {
            A11 = -conditioned_base[i](2) * conditioned_attach[i].t();
            A13 = conditioned_base[i](0) * conditioned_attach[i].t();
            A22 = -conditioned_base[i](2) * conditioned_attach[i].t();
            A23 = conditioned_base[i](1) * conditioned_attach[i].t();
            A(cv::Rect(0, i * 2, 3, 1)) << A11(0, 0), A11(0, 1), A11(0, 2);
            A(cv::Rect(6, i * 2, 3, 1)) << A13(0, 0), A13(0, 1), A13(0, 2);
            A(cv::Rect(3, i * 2 + 1, 3, 1)) << A22(0, 0), A22(0, 1), A22(0, 2);
            A(cv::Rect(6, i * 2 + 1, 3, 1)) << A23(0, 0), A23(0, 1), A23(0, 2);
        }

        return A;
    }

    cv::Matx33f solve_dlt_homography2D(const cv::Mat_<float> &A) {
        cv::Matx33f Hhat;
        cv::Mat_<float> U, W, VT;
        cv::SVD::compute(A, W, U, VT, 4);

        for (int i = 0; i < 3; i++) {
            for (int j = 0; j < 3; j++) {
                Hhat(i, j) = VT(VT.rows - 1, i * 3 + j);
            }
        }

        return Hhat;
    }


    cv::Matx33f decondition_homography2D(const cv::Matx33f &T_base, const cv::Matx33f &T_attach, const cv::Matx33f &H) {
        cv::Matx33f Hout;
        Hout = T_base.inv() * H * T_attach;
        return Hout;
    }

    cv::Matx33f homography2D(const std::vector<cv::Vec3f> &base, const std::vector<cv::Vec3f> &attach) {
        std::vector<cv::Vec3f> conditioned_base;
        std::vector<cv::Vec3f> conditioned_attach;
        cv::Matx33f T_base;
        cv::Matx33f T_attach;
        cv::Mat_<float> A;
        cv::Matx33f Hhat;
        cv::Matx33f H;
        T_base = getCondition2D(base);
        T_attach = getCondition2D(attach);

        for (const auto &i: base) {
            conditioned_base.push_back(T_base * i);
        }
        for (const auto &i: attach) {
            conditioned_attach.push_back(T_attach * i);
        }

        A = getDesignMatrix_homography2D(conditioned_base, conditioned_attach);
        Hhat = solve_dlt_homography2D(A);
        H = decondition_homography2D(T_base, T_attach, Hhat);
        return H;
    }

}
